// 文件：src/LangGraph/instance/retrievers.ts
import { PGVectorStore } from '@langchain/community/vectorstores/pgvector';
import { OllamaEmbeddings } from '@langchain/ollama';
import { Document } from '@langchain/core/documents';

// 检索结果类型
export type RetrieverResult = Array<{
  id: string;
  text: string;
  meta: any;
  score?: number;
}>;

// 通用向量存储初始化配置
const getEmbeddings = () => new OllamaEmbeddings({
  baseUrl: 'http://192.168.31.159:11434',
  model: 'quentinz/bge-large-zh-v1.5:latest',
  requestOptions: {
    useMmap: true,
    numThread: 10,
    numGpu: 1,
  },
});

// 通用PostgreSQL连接配置（可从环境变量读取）
const getPostgresConfig = () => ({
  host: process.env.DB_HOST || '118.178.106.254',
  port: parseInt(process.env.DB_PORT || '15432'),
  user: process.env.DB_USER || 'postgres',
  password: process.env.DB_PASSWORD || '6a29c7071ae1f1b9',
  database: process.env.DB_NAME || 'postgres',
});

/**
 * TopK检索器 - 基于相似度分数返回前K个文档
 */
export async function topKRetriever(
  collectionName: string,
  k: number = 5
): Promise<(query: string) => Promise<RetrieverResult>> {
  const vectorStore = await PGVectorStore.initialize(
    getEmbeddings(),
    {
      tableName: `embedding_${collectionName}_1024`,
      columns: {
        idColumnName: 'embedding_id',
        vectorColumnName: 'embedding',
        contentColumnName: 'text',
        metadataColumnName: 'metadata',
      },
      postgresConnectionOptions: getPostgresConfig(),
    }
  );

  const retriever = vectorStore.asRetriever(k);

  return async (query: string): Promise<RetrieverResult> => {
    const docs = await retriever.invoke(query);
    return docs.map((doc: Document, index: number) => ({
      id: doc.metadata?.embedding_id || `doc_${index}`,
      text: doc.pageContent,
      meta: {
        ...doc.metadata,
        source: doc.metadata?.source || collectionName,
        chunkIndex: doc.metadata?.chunkIndex || index,
      },
      score: doc.metadata?.score,
    }));
  };
}

/**
 * 混合检索器 - 结合语义检索和关键词检索（简化版：使用MMR多样性检索）
 */
export async function hybridRetriever(
  collectionName: string,
  k: number = 5,
  fetchK: number = 20
): Promise<(query: string) => Promise<RetrieverResult>> {
  const vectorStore = await PGVectorStore.initialize(
    getEmbeddings(),
    {
      tableName: `embedding_${collectionName}_1024`,
      columns: {
        idColumnName: 'embedding_id',
        vectorColumnName: 'embedding',
        contentColumnName: 'text',
        metadataColumnName: 'metadata',
      },
      postgresConnectionOptions: getPostgresConfig(),
    }
  );

  // 使用MMR（最大边际相关性）来增加结果多样性
  const retriever = vectorStore.asRetriever({
    k,
    fetchK,
    searchType: 'mmr',
    mmrThreshold: 0.5, // 多样性阈值
  });

  return async (query: string): Promise<RetrieverResult> => {
    const docs = await retriever.invoke(query);
    return docs.map((doc: Document, index: number) => ({
      id: doc.metadata?.embedding_id || `doc_${index}`,
      text: doc.pageContent,
      meta: {
        ...doc.metadata,
        source: doc.metadata?.source || collectionName,
        chunkIndex: doc.metadata?.chunkIndex || index,
      },
      score: doc.metadata?.score,
    }));
  };
}

/**
 * 时间感知检索器 - 优先返回最近的文档
 */
export async function timeAwareRetriever(
  collectionName: string,
  k: number = 5,
  timeWeight: number = 0.3
): Promise<(query: string) => Promise<RetrieverResult>> {
  const vectorStore = await PGVectorStore.initialize(
    getEmbeddings(),
    {
      tableName: `embedding_${collectionName}_1024`,
      columns: {
        idColumnName: 'embedding_id',
        vectorColumnName: 'embedding',
        contentColumnName: 'text',
        metadataColumnName: 'metadata',
      },
      postgresConnectionOptions: getPostgresConfig(),
    }
  );

  const retriever = vectorStore.asRetriever(k * 2); // 多检索一些，再按时间排序

  return async (query: string): Promise<RetrieverResult> => {
    const docs = await retriever.invoke(query);
    const now = Date.now();
    
    // 根据时间戳和相似度分数计算加权分数
    const scored = docs.map((doc: Document, index: number) => {
      const docTime = doc.metadata?.Date 
        ? (typeof doc.metadata.Date === 'string' 
            ? new Date(doc.metadata.Date).getTime() 
            : doc.metadata.Date.getTime())
        : 0;
      const timeScore = docTime > 0 ? (now - docTime) / (365 * 24 * 60 * 60 * 1000) : 1; // 归一化到年
      const similarityScore = doc.metadata?.score || (1 - index / docs.length);
      const finalScore = similarityScore * (1 - timeWeight) + (1 - timeScore) * timeWeight;
      
      return {
        id: doc.metadata?.embedding_id || `doc_${index}`,
        text: doc.pageContent,
        meta: {
          ...doc.metadata,
          source: doc.metadata?.source || collectionName,
          chunkIndex: doc.metadata?.chunkIndex || index,
        },
        score: finalScore,
        _originalScore: similarityScore,
        _timeScore: timeScore,
      };
    });

    // 按最终分数排序并返回前k个
    return scored
      .sort((a, b) => (b.score || 0) - (a.score || 0))
      .slice(0, k)
      .map(({ _originalScore, _timeScore, ...rest }) => rest);
  };
}

/**
 * 用户感知检索器 - 根据用户信息（部门、角色等）过滤和排序文档
 */
export async function userAwareRetriever(
  collectionName: string,
  userInfo: { id: string; dept?: string; role?: string; [key: string]: any },
  k: number = 5
): Promise<(query: string) => Promise<RetrieverResult>> {
  const vectorStore = await PGVectorStore.initialize(
    getEmbeddings(),
    {
      tableName: `embedding_${collectionName}_1024`,
      columns: {
        idColumnName: 'embedding_id',
        vectorColumnName: 'embedding',
        contentColumnName: 'text',
        metadataColumnName: 'metadata',
      },
      postgresConnectionOptions: getPostgresConfig(),
    }
  );

  const retriever = vectorStore.asRetriever(k * 2); // 多检索一些，再按用户信息过滤

  return async (query: string): Promise<RetrieverResult> => {
    const docs = await retriever.invoke(query);
    
    // 根据用户信息对文档进行评分和过滤
    const scored = docs.map((doc: Document, index: number) => {
      const meta = doc.metadata || {};
      let userRelevance = 0.5; // 默认相关性

      // 如果文档有部门信息，匹配用户部门
      if (userInfo.dept && meta.department) {
        if (meta.department === userInfo.dept) {
          userRelevance += 0.3;
        }
      }

      // 如果文档有角色信息，匹配用户角色
      if (userInfo.role && meta.role) {
        if (meta.role === userInfo.role) {
          userRelevance += 0.2;
        }
      }

      // 如果文档标记为用户特定
      if (meta.userId && meta.userId === userInfo.id) {
        userRelevance += 0.5;
      }

      const similarityScore = meta.score || (1 - index / docs.length);
      const finalScore = similarityScore * 0.7 + userRelevance * 0.3;

      return {
        id: meta.embedding_id || `doc_${index}`,
        text: doc.pageContent,
        meta: {
          ...meta,
          source: meta.source || collectionName,
          chunkIndex: meta.chunkIndex || index,
          userInfo, // 保存用户信息用于后续处理
        },
        score: finalScore,
      };
    });

    // 按最终分数排序并返回前k个
    return scored
      .sort((a, b) => (b.score || 0) - (a.score || 0))
      .slice(0, k);
  };
}
